""" 
author: Lucywu & cursor IDE
date: 2025-03-26
version: 1.0
"""
import os
import sys
import fitz  # PyMuPDF
import shutil
from typing import List, Dict
import re
import time

class PDFProcessor:
    def __init__(self, input_path: str):
        self.input_path = input_path
        self.output_path = self._generate_output_path()
        
    def _generate_output_path(self) -> str:
        """Generate output path by adding '_new' before extension"""
        base, ext = os.path.splitext(self.input_path)
        return f"{base}_new{ext}"
    
    def validate_input(self) -> bool:
        """
        Validate input file and permissions
        验证输入文件和权限
        """
        print("正在验证输入文件... / Validating input file...")
        
        if not os.path.exists(self.input_path):
            print("错误：输入文件不存在 / Error: Input file does not exist")
            return False
            
        if not self.input_path.lower().endswith('.pdf'):
            print("错误：输入文件不是PDF格式 / Error: Input file is not a PDF")
            return False
            
        output_dir = os.path.dirname(self.output_path)
        if not os.access(output_dir, os.W_OK):
            print("错误：没有写入权限 / Error: No write permission in target directory")
            return False
            
        try:
            # Try to open PDF to verify it's valid
            with fitz.open(self.input_path) as doc:
                pass
        except Exception as e:
            print(f"错误：无效的PDF文件 / Error: Invalid PDF file - {str(e)}")
            return False
            
        return True
        
    def analyze_content(self, doc: fitz.Document) -> List[Dict]:
        """
        Analyze document content and generate bookmark structure
        分析文档内容并生成书签结构
        """
        print("正在分析文档内容... / Analyzing document content...")
        bookmarks = []
        added_titles = set()
        
        # 获取第一页的页边距作为参考
        first_page = doc[0]
        mediabox = first_page.mediabox  # 媒体框
        cropbox = first_page.cropbox    # 裁剪框
        
        # 计算实际的页边距
        margin_top = cropbox.y0 - mediabox.y0
        margin_bottom = mediabox.y1 - cropbox.y1
        
        # 如果页边距无效（为0或负值），使用默认值
        if margin_top <= 0:
            margin_top = 50  # 默认上边距50点
        if margin_bottom <= 0:
            margin_bottom = 50  # 默认下边距50点
        
        # 定义标题模式和对应的层级
        title_patterns = [
            # 第一级：中文数字编号
            (r'^[一二三四五六七八九十]+、.*$', 1),  # 一、标题内容
            (r'^第[一二三四五六七八九十]+[章节篇].*$', 1),  # 第一章标题内容
            
            # 第二级：带括号的中文数字
            (r'^[（\(][一二三四五六七八九十]+[）\)].*$', 2),  # （一）标题内容
            (r'^[（\(][一二三四五六七八九十]+[）\)]\d+\..*$', 2),  # （一）1.标题内容
            (r'^[（\(][一二三四五六七八九十]+[）\)]\d+.*$', 2),  # （一）1标题内容
            
            # 第三级：阿拉伯数字带点
            (r'^\d+\.\s*.*$', 3),  # 1. 标题内容
            (r'^[（\(][一二三四五六七八九十]+[）\)]\d+\.\d+.*$', 3),  # （一）1.1标题内容
            
            # 第四级：带括号的阿拉伯数字
            (r'^[（\(]\d+[）\)].*$', 4),  # （1）标题内容
            (r'^\d+\.\d+\.\d+.*$', 4),  # 1.1.1标题内容
            
            # 第五级：带右括号的阿拉伯数字
            (r'^\d+[）\)].*$', 5),  # 1）标题内容
            (r'^\d+\.\d+\.\d+\.\d+.*$', 5),  # 1.1.1.1标题内容
            
            # 特殊标题（一级）- 这些可以单独成行
            (r'^(引言|前言|摘要|结论|参考文献|附件\d*|附录\d*)$', 1),
        ]
        
        # 编译正则表达式
        compiled_patterns = [(re.compile(pattern), level) for pattern, level in title_patterns]
        
        # 获取正文主要字体大小
        font_sizes = {}
        for page_num in range(len(doc)):
            page = doc[page_num]
            blocks = page.get_text("dict")["blocks"]
            
            for block in blocks:
                if "lines" not in block or not block["lines"]:
                    continue
                for line in block["lines"]:
                    for span in line["spans"]:
                        size = round(span["size"], 1)
                        font_sizes[size] = font_sizes.get(size, 0) + 1
        
        body_text_size = max(font_sizes.items(), key=lambda x: x[1])[0]  # 最常见字号
        
        # 扫描文档识别标题
        for page_num in range(len(doc)):
            page = doc[page_num]
            blocks = sorted(page.get_text("dict")["blocks"], key=lambda b: (b["bbox"][1], b["bbox"][0]))
            
            for block in blocks:
                if "lines" not in block or not block["lines"]:
                    continue
                
                # 只处理单行或两行的段落
                if len(block["lines"]) > 2:
                    continue
                
                # 只获取第一行的文本作为潜在标题
                first_line = block["lines"][0]
                line_text = "".join(span["text"] for span in first_line["spans"]).strip()
                if not line_text:  # 跳过空行
                    continue
                
                # 获取字体大小
                max_font_size = max(span["size"] for span in first_line["spans"])
                
                if not line_text or len(line_text) > 100:  # 跳过空文本或过长的文本
                    continue
                
                # 跳过页眉页脚
                y_pos = first_line["bbox"][1]
                y_bottom = first_line["bbox"][3]
                if y_pos < margin_top or y_bottom > (page.rect.height - margin_bottom):
                    continue
                
                # 检查标题模式
                matched_level = None
                for pattern, level in compiled_patterns:  # 使用编译后的模式
                    if pattern.match(line_text):
                        matched_level = level
                        break
                
                # 判断是否为标题
                is_title = (
                    matched_level and  # 匹配标题格式
                    max_font_size >= body_text_size * 0.9 and  # 字体不小于正文
                    line_text not in added_titles and  # 避免重复
                    len(line_text.strip()) >= 2  # 标题至少2个字符
                )
                
                if is_title:
                    # 如果标题中包含句号，只取句号前的部分
                    title_text = line_text.split('。')[0]
                    # 获取文本块的精确位置
                    bookmarks.append({
                        "title": title_text,
                        "page": page_num,
                        "level": matched_level,
                        "y": y_pos  # 记录Y坐标
                    })
                    added_titles.add(line_text)
        
        return bookmarks
        
    def process_pdf(self):
        """
        Main processing function
        主要处理函数
        """
        if not self.validate_input():
            return False
            
        doc = None
        src_doc = None
        temp_path = None
        success = False
        max_retries = 3
        
        try:
            # Create a temporary file path
            temp_dir = os.path.dirname(self.output_path)
            temp_path = os.path.join(temp_dir, f"temp_{os.urandom(4).hex()}.pdf")
            
            print("正在创建PDF副本... / Creating PDF copy...")
            # Open source document
            src_doc = fitz.open(self.input_path)
            
            # Create new document
            doc = fitz.open()
            
            # Copy pages from source to new document
            doc.insert_pdf(src_doc)
            src_doc.close()
            src_doc = None
            
            # Remove existing bookmarks
            print("正在移除现有书签... / Removing existing bookmarks...")
            doc.set_toc([])
            
            # Analyze content and generate new bookmarks
            bookmarks = self.analyze_content(doc)
            
            if bookmarks:
                print(f"找到 {len(bookmarks)} 个潜在的书签 / Found {len(bookmarks)} potential bookmarks")
                
                # 按页码排序
                bookmarks.sort(key=lambda x: x["page"])
                
                # 规范化书签层级
                normalized_bookmarks = []
                current_levels = {}  # 记录当前各层级的状态
                
                # 确保第一个书签是一级标题
                if bookmarks:
                    first_bookmark = bookmarks[0]
                    normalized_bookmarks.append([1, first_bookmark["title"], first_bookmark["page"] + 1, first_bookmark["y"]])
                    current_levels[1] = first_bookmark["title"]
                
                # 处理剩余书签
                for bm in bookmarks[1:]:
                    title = bm["title"]
                    page = bm["page"]
                    level = bm["level"]
                    y_pos = bm["y"]
                    
                    # 获取当前最深层级
                    max_current_level = max(current_levels.keys()) if current_levels else 0
                    
                    # 调整层级
                    if level > max_current_level + 1:
                        level = max_current_level + 1
                    elif level < 1:
                        level = 1
                    
                    # 更新层级状态
                    current_levels[level] = title
                    deeper_levels = [l for l in list(current_levels.keys()) if l > level]
                    for l in deeper_levels:
                        current_levels.pop(l)
                    
                    # 添加带有Y坐标的书签
                    normalized_bookmarks.append([level, title, page + 1, y_pos])
                
                # Apply new bookmarks
                print("正在应用新书签... / Applying new bookmarks...")
                doc.set_toc(normalized_bookmarks)
            
            # Save to temporary file first
            doc.save(temp_path, garbage=4, deflate=True)
            doc.close()
            doc = None
            
            # 尝试替换文件，如果失败则重试
            for attempt in range(max_retries):
                try:
                    if os.path.exists(self.output_path):
                        os.remove(self.output_path)
                    os.rename(temp_path, self.output_path)
                    break
                except PermissionError:
                    if attempt < max_retries - 1:
                        print(f"文件被占用，等待重试... / File is in use, waiting to retry... ({attempt + 1}/{max_retries})")
                        time.sleep(2)  # 等待2秒后重试
                    else:
                        raise Exception("文件被占用，无法完成操作 / File is in use, cannot complete operation")
            
            print(f"处理完成！新文件保存为：{self.output_path}")
            print(f"Processing complete! New file saved as: {self.output_path}")
            success = True
            return True
            
        except Exception as e:
            print(f"处理过程中出错 / Error during processing: {str(e)}")
            return False
            
        finally:
            # Close documents if still open
            if src_doc:
                try:
                    src_doc.close()
                except:
                    pass
                    
            if doc:
                try:
                    doc.close()
                except:
                    pass
                    
            # Clean up temporary file if it exists and process failed
            if not success and temp_path and os.path.exists(temp_path):
                try:
                    os.remove(temp_path)
                except Exception as e:
                    print(f"清理临时文件失败 / Failed to clean up temporary file: {str(e)}")

def main():
    if len(sys.argv) != 2:
        print("用法 / Usage: python pdf_processor.py /path/to/input.pdf")
        sys.exit(1)
        
    processor = PDFProcessor(sys.argv[1])
    success = processor.process_pdf()
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()
